import autolens as al
import autolens.plot as aplt
import autofit as af
import numpy as np
from matplotlib import pyplot as plt
from astropy.io import fits
from os import path

'''necessary parameters'''
###source and lens redshift
redshift_l = 1.7
redshift_s = 2.5579

###data profile
pixel_scales = 0.025
pixel_scale_precision = 0.005
position_noise = [pixel_scales]*4
noise = [1]*4

###the position and flux of the four point sources, to accelerate the fitting
positions = np.array([[-0.327,-0.0057*15],[-0.152,-0.0550*15],[0.723,-0.0295*15],[0.398,0.0283*15]])-np.array([[0.243,-0.0108*15]])
fluxes = [34.6633,30.1659,21.1212,21.1245]

###the folder that output the results
dataname = "F814W_coadd"
name="overview_point_source_t33"

'''run the fitting'''
###the sample grid
grid = al.Grid2D.uniform(shape_native=(500, 500), pixel_scales=pixel_scales)
solver = al.PositionsSolver(
    grid=grid,
    pixel_scale_precision=pixel_scale_precision,
)

###input the data image
image = al.Array2D.from_fits(
    file_path=path.join("./data/F814W_coadd_cut_f5.fits"), pixel_scales=0.025)

###input the relative position of the point sources
point_dataset = al.PointDataset(
    name="point_0",
    positions=al.Grid2DIrregular(positions),
    positions_noise_map=al.ValuesIrregular(position_noise),
    fluxes_noise_map=al.ValuesIrregular(noise),
    fluxes=al.ValuesIrregular(fluxes)
)
point_dict = al.PointDict(point_dataset_list=[point_dataset])

'''the model and the fitting'''
###lens and source model
lens_galaxy_model = af.Model(al.Galaxy, redshift=redshift_l, mass=al.mp.EllIsothermal, shear=al.mp.ExternalShear)
source_galaxy_model = af.Model(al.Galaxy, redshift=redshift_s, point_0=al.ps.Point)

galaxies = af.Collection(lens=lens_galaxy_model, source=source_galaxy_model)
model = af.Collection(galaxies=galaxies)

###mcmc fitting
search = af.DynestyStatic(name=name,nlive=200,number_of_cores=150,walks=10)

analysis = al.AnalysisPoint(point_dict=point_dict, solver=solver)

result = search.fit(model=model, analysis=analysis)

fit = al.FitPointDataset(
    point_dataset=point_dict["point_0"],
    tracer=result.max_log_likelihood_tracer,
    positions_solver=solver)




